{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import os\n",
    "import matplotlib.pyplot as plt\n",
    "import cv2\n",
    "import re\n",
    "import warnings\n",
    "from scipy.sparse import spdiags\n",
    "from scipy import signal\n",
    "from scipy.signal import butter\n",
    "import scipy\n",
    "\n",
    "%matplotlib inline\n",
    "\n",
    "# Try importing ipywidgets\n",
    "try:\n",
    "    import ipywidgets as widgets\n",
    "    from IPython.display import display, clear_output\n",
    "except ImportError as e:\n",
    "    widgets = None\n",
    "    print(\"ipywidgets is not available in this environment. Interactive features will be disabled.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# FILL THIS IN! This is the path to your preprocessed dataset folder.\n",
    "preprocessed_dataset_path = '/your/preprocessed/dataset/path/here'\n",
    "\n",
    "input_files = [filename for filename in os.listdir(preprocessed_dataset_path) if filename.endswith(\".npy\") and re.match(r\".*input\\d+\\.npy\", filename)]\n",
    "label_files = [input_file.replace(\"input\", \"label\") for input_file in input_files if input_file.replace(\"input\", \"label\") in os.listdir(preprocessed_dataset_path)]\n",
    "\n",
    "if len(input_files) == 0 or len(label_files) == 0:\n",
    "    print(\"Error: No valid input or label files found in the directory.\")\n",
    "else:\n",
    "    print(f\"{len(input_files)} input files and {len(label_files)} label files detected.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Visualizer:\n",
    "    def __init__(self, preprocessed_dataset_path, input_files, label_files):\n",
    "        self.input_files = sorted(input_files)\n",
    "        self.label_files = sorted(label_files)\n",
    "        self.chunk_dropdown = None\n",
    "        self.fig = None\n",
    "        self.ax = None\n",
    "        self.frame_index_slider = None\n",
    "        self.current_chunk_index = 0\n",
    "        self.current_frame_index = 0\n",
    "        self.dataset_path = preprocessed_dataset_path\n",
    "        self.avg_heart_rate = None\n",
    "\n",
    "        self.diff_flag = \"DiffNormalized\" in preprocessed_dataset_path\n",
    "\n",
    "        # Only initialize widgets if the library is available\n",
    "        if widgets:\n",
    "            self.create_visualization()\n",
    "        else:\n",
    "            print(\"Running in a non-widget environment. Visualization will be non-interactive.\")\n",
    "\n",
    "    def create_visualization(self):\n",
    "        self.create_dropdown()\n",
    "        self.create_figure()\n",
    "        self.create_slider()\n",
    "\n",
    "    def create_dropdown(self):\n",
    "        # Check if widgets are available\n",
    "        if widgets:\n",
    "            self.chunk_dropdown = widgets.Dropdown(\n",
    "                options=self.input_files,\n",
    "                description='Chunk:',\n",
    "                layout=widgets.Layout(width='800px'),\n",
    "                value=None\n",
    "            )\n",
    "            self.chunk_dropdown.observe(self.chunk_selection_changed, names='value')\n",
    "\n",
    "            display(self.chunk_dropdown)\n",
    "        else:\n",
    "            print(\"Dropdown cannot be displayed: ipywidgets not available.\")\n",
    "\n",
    "    def create_figure(self):\n",
    "        self.fig, self.ax = plt.subplots(figsize=(8, 4))\n",
    "        self.ax.set_xticks([])\n",
    "        self.ax.set_yticks([])\n",
    "        plt.tight_layout()\n",
    "\n",
    "    def create_slider(self):\n",
    "        if widgets:\n",
    "            num_frames = self.get_num_frames()\n",
    "            self.frame_index_slider = widgets.IntSlider(\n",
    "                min=0,\n",
    "                max=num_frames - 1,\n",
    "                step=1,\n",
    "                value=0,\n",
    "                description='Frame:',\n",
    "                layout=widgets.Layout(width='800px')\n",
    "            )\n",
    "            self.frame_index_slider.observe(self.frame_index_changed, names='value')\n",
    "\n",
    "            display(self.frame_index_slider)\n",
    "        else:\n",
    "            print(\"Slider cannot be displayed: ipywidgets not available.\")\n",
    "\n",
    "    def chunk_selection_changed(self, change):\n",
    "        selected_chunk = change.new\n",
    "        self.current_chunk_index = self.input_files.index(selected_chunk)\n",
    "        self.current_frame_index = 0\n",
    "        if self.frame_index_slider:\n",
    "            self.frame_index_slider.value = 0\n",
    "        self.update_frame()\n",
    "\n",
    "    def frame_index_changed(self, change):\n",
    "        self.current_frame_index = change.new\n",
    "        self.update_frame()\n",
    "\n",
    "    def get_num_frames(self):\n",
    "        input_file = self.input_files[self.current_chunk_index]\n",
    "        return np.load(os.path.join(self.dataset_path, input_file)).shape[0]\n",
    "\n",
    "    def load_frame(self):\n",
    "        input_file = self.input_files[self.current_chunk_index]\n",
    "        input_data = np.load(os.path.join(self.dataset_path, input_file))\n",
    "        print(f'The shape of the loaded chunk is {np.shape(input_data)}.')\n",
    "\n",
    "        if input_data.shape[-1] == 6:\n",
    "            rgb_frames = input_data[..., 3:]\n",
    "            diff_normalized_frames = input_data[..., :3]\n",
    "\n",
    "            # Create a side-by-side visualization\n",
    "            frame = np.concatenate((rgb_frames, diff_normalized_frames), axis=2)\n",
    "        elif input_data.shape[-1] == 3:\n",
    "            # Use RGB frames directly for visualization\n",
    "            frame = input_data\n",
    "        else:\n",
    "            raise ValueError(\"Invalid input_data shape. Expected shape (..., 3) or (..., 6).\")\n",
    "\n",
    "        return frame[self.current_frame_index]\n",
    "\n",
    "    def load_label_frame(self):\n",
    "        \"\"\"\n",
    "        Load the corresponding label data for the currently selected chunk.\n",
    "        This method assumes the labels are stored in .npy files and are aligned\n",
    "        with the input files in the same folder.\n",
    "        \"\"\"\n",
    "        label_file = self.label_files[self.current_chunk_index]\n",
    "        label_data = np.load(os.path.join(self.dataset_path, label_file))\n",
    "        print(f'The shape of the loaded label file is {np.shape(label_data)}.')\n",
    "\n",
    "        # Process the signal as needed (e.g., detrending, filtering)\n",
    "        label_data = self._process_signal(label_data)\n",
    "\n",
    "        return label_data\n",
    "\n",
    "    def _process_signal(self, signal):\n",
    "        \"\"\"\n",
    "        A helper function to process the signal (e.g., detrend and filter the PPG signal).\n",
    "        You can adapt this based on your signal processing needs.\n",
    "        \"\"\"\n",
    "        fs = 30  # Sampling frequency\n",
    "        use_bandpass = True\n",
    "\n",
    "        # Detrend signal or its cumulative sum based on diff_flag\n",
    "        if self.diff_flag:\n",
    "            # If the predictions and labels are the 1st derivative of the PPG signal\n",
    "            detrended_signal = self._detrend(np.cumsum(signal), 100)\n",
    "        else:\n",
    "            # Detrend the original signal\n",
    "            detrended_signal = self._detrend(signal, 100)\n",
    "\n",
    "        # Apply bandpass filter if needed\n",
    "        if use_bandpass:\n",
    "            # Bandpass filter between [0.75, 2.5] Hz (45-150 bpm)\n",
    "            b, a = butter(1, [0.75 / fs * 2, 2.5 / fs * 2], btype='bandpass')\n",
    "            filtered_signal = scipy.signal.filtfilt(b, a, np.double(detrended_signal))\n",
    "            return filtered_signal\n",
    "\n",
    "        return detrended_signal\n",
    "\n",
    "    def _detrend(self, input_signal, lambda_value):\n",
    "        \"\"\"\n",
    "        Detrend the input signal.\n",
    "        This method removes long-term trends from the signal using a smoothing filter.\n",
    "        \"\"\"\n",
    "        signal_length = input_signal.shape[0]\n",
    "        H = np.identity(signal_length)\n",
    "        ones = np.ones(signal_length)\n",
    "        minus_twos = -2 * np.ones(signal_length)\n",
    "        diags_data = np.array([ones, minus_twos, ones])\n",
    "        diags_index = np.array([0, 1, 2])\n",
    "        D = spdiags(diags_data, diags_index, signal_length - 2, signal_length).toarray()\n",
    "\n",
    "        # Apply smoothing\n",
    "        detrended_signal = np.dot(\n",
    "            (H - np.linalg.inv(H + (lambda_value ** 2) * np.dot(D.T, D))), input_signal\n",
    "        )\n",
    "        return detrended_signal\n",
    "\n",
    "    def create_label_figure(self):\n",
    "        # Create figure with two subplots: time-domain and frequency-domain\n",
    "        self.label_fig, (self.label_ax_time, self.label_ax_fft) = plt.subplots(2, 1, figsize=(8, 6))\n",
    "        self.label_ax_time.set_xlabel('Frame #')\n",
    "        self.label_ax_time.set_ylabel('Magnitude')\n",
    "        self.label_ax_fft.set_xlabel('Frequency')\n",
    "        self.label_ax_fft.set_ylabel('Magnitude')\n",
    "\n",
    "    def update_label_frame(self):\n",
    "        def _next_power_of_2(x):\n",
    "            \"\"\"Calculate the nearest power of 2.\"\"\"\n",
    "            return 1 if x == 0 else 2 ** (x - 1).bit_length()\n",
    "\n",
    "        clear_output(wait=True)\n",
    "\n",
    "        label_frame = self.load_label_frame()\n",
    "\n",
    "        # Zero-padding the signal to improve frequency resolution\n",
    "        N = _next_power_of_2(label_frame.shape[0])\n",
    "        fs = 30  # Sampling frequency\n",
    "        ppg_label_fft = label_frame\n",
    "        \n",
    "        ppg_label_f, ppg_label_pxx = signal.periodogram(ppg_label_fft, fs=fs, nfft=N, detrend=False)\n",
    "\n",
    "        # Find the frequency corresponding to the highest power within the heart rate range (0.75 - 2.5 Hz)\n",
    "        heart_rate_range = (0.75, 2.5)\n",
    "        idx = np.where((ppg_label_f >= heart_rate_range[0]) & (ppg_label_f <= heart_rate_range[1]))[0]\n",
    "        dominant_frequency = ppg_label_f[idx[np.argmax(ppg_label_pxx[idx])]]\n",
    "        self.avg_heart_rate = dominant_frequency * 60  # Convert Hz to bpm (beats per minute)\n",
    "\n",
    "        # Clear the figure to prevent double visualization\n",
    "        self.label_fig.clear()\n",
    "\n",
    "        # Create the time-domain plot\n",
    "        self.label_ax_time = self.label_fig.add_subplot(211)  # First subplot for time-domain\n",
    "        self.label_ax_time.plot(label_frame)\n",
    "        self.label_ax_time.set_xlabel('Frame #')\n",
    "        self.label_ax_time.set_ylabel('Magnitude')\n",
    "\n",
    "        # Create the frequency-domain plot\n",
    "        self.label_ax_fft = self.label_fig.add_subplot(212)  # Second subplot for frequency-domain\n",
    "        self.label_ax_fft.plot(ppg_label_f, ppg_label_pxx / ppg_label_pxx.max())\n",
    "        self.label_ax_fft.set_xlabel('Frequency (Hz)')\n",
    "        self.label_ax_fft.set_ylabel('Magnitude')\n",
    "        self.label_ax_fft.set_xlim([0, 5])  # Limit the x-axis to relevant frequencies (0-5 Hz)\n",
    "\n",
    "        self.label_fig.suptitle(self.label_files[self.current_chunk_index], y=1.02)\n",
    "        self.label_fig.tight_layout()\n",
    "\n",
    "        # Display the average heart rate as standalone text below the entire figure\n",
    "        self.label_fig.text(0.5, -0.05, f'Average Heart Rate: {self.avg_heart_rate:.2f} bpm',\n",
    "                            ha='center', va='center', fontsize=16, color='red', weight='bold')\n",
    "\n",
    "    def update_frame(self):\n",
    "        clear_output(wait=True)\n",
    "\n",
    "        self.fig.clear()\n",
    "\n",
    "        frame = self.load_frame()\n",
    "\n",
    "        # Recreate the axes for the figure and draw the image\n",
    "        self.ax = self.fig.add_subplot(111)\n",
    "        self.ax.imshow(frame)\n",
    "\n",
    "        self.ax.set_xticks([])\n",
    "        self.ax.set_yticks([])\n",
    "\n",
    "        self.ax.set_title(f\"{self.input_files[self.current_chunk_index]} - Frame {self.current_frame_index}\")\n",
    "\n",
    "        self.fig.tight_layout()\n",
    "\n",
    "        self.update_label_frame()\n",
    "\n",
    "        display(self.chunk_dropdown, self.frame_index_slider, self.fig)\n",
    "\n",
    "    def update_frame(self):\n",
    "        clear_output(wait=True)\n",
    "\n",
    "        self.fig.clear()  # This ensures the whole figure is cleared before re-drawing\n",
    "\n",
    "        frame = self.load_frame()\n",
    "\n",
    "        # Recreate the axes for the figure and draw the image\n",
    "        self.ax = self.fig.add_subplot(111)\n",
    "        self.ax.imshow(frame)\n",
    "\n",
    "        self.ax.set_xticks([])\n",
    "        self.ax.set_yticks([])\n",
    "\n",
    "        self.ax.set_title(f\"{self.input_files[self.current_chunk_index]} - Frame {self.current_frame_index}\")\n",
    "\n",
    "        self.fig.tight_layout()\n",
    "\n",
    "        self.update_label_frame()\n",
    "\n",
    "        display(self.chunk_dropdown, self.frame_index_slider, self.fig, self.label_fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the visualizer instance\n",
    "visualizer = Visualizer(preprocessed_dataset_path, input_files, label_files)\n",
    "\n",
    "# Create the label figure\n",
    "visualizer.create_label_figure()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "rppg-toolbox",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.20"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
